import socket
import json
import requests

# Ollama服务地址
OLLAMA_BASE_URL = "http://localhost:11434"

def generate_response(prompt, model="llama2"):
    """调用Ollama API生成回复"""
    try:
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "stream": False
            }
        )
        response.raise_for_status()
        return response.json()["response"].strip()
    except Exception as e:
        print(f"生成回复时出错: {e}")
        return "抱歉，生成回复时出错。"

def start_client():
    """启动TCP客户端"""
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect(('localhost', 9999))
    
    try:
        # 接收服务端初始消息
        server_message = client.recv(4096).decode()
        print(f"服务端: {server_message}")
        
        while True:
            # 生成客户端回复
            client_response = generate_response(server_message)
            print(f"客户端回复: {client_response}")
            
            # 发送回复给服务端
            client.send(client_response.encode())
            
            # 接收服务端新消息
            server_message = client.recv(4096).decode()
            if not server_message:
                break
            print(f"服务端: {server_message}")
            
    except Exception as e:
        print(f"客户端通信时出错: {e}")
    finally:
        client.close()
        print("客户端已关闭")

if __name__ == "__main__":
    start_client()
